{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 对话记忆 Conversations with Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade langchain\n",
    "%pip install python-dotenv\n",
    "%pip install openai\n",
    "%pip -q install huggingface_hub transformers\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 导入和连接 OPENAI"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基本感知对话缓冲记忆 ConversationBufferMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "from langchain import OpenAI\n",
    "import json\n",
    "\n",
    "# load_dotenv('.env.local')\n",
    "# openai.api_key = os.environ.get(\"OPENAI_API_TOKEN\")\n",
    "\n",
    "OPENAI_API_TOKEN=\"\"\n",
    "OpenAI.openai_api_key = OPENAI_API_TOKEN\n",
    "print(f\"{OpenAI.openai_api_key}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.conversation.memory import ConversationBufferMemory\n",
    "\n",
    "from langchain.chains import ConversationChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model_name='text-davinci-003', \n",
    "             temperature=0, \n",
    "             max_tokens = 256,\n",
    "             openai_api_key=OPENAI_API_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationBufferMemory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    memory=memory\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi there! I am Sam\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\""
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: How are you today?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" I'm doing great, thanks for asking! How about you?\""
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"How are you today?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: How are you today?\n",
      "AI:  I'm doing great, thanks for asking! How about you?\n",
      "Human: I'm good thank you. Can you help me with some customer support?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Absolutely! What kind of customer support do you need?'"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"I'm good thank you. Can you help me with some customer support?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: How are you today?\n",
      "AI:  I'm doing great, thanks for asking! How about you?\n",
      "Human: I'm good thank you. Can you help me with some customer support?\n",
      "AI:  Absolutely! What kind of customer support do you need?\n"
     ]
    }
   ],
   "source": [
    "print(conversation.memory.buffer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以上对话，看起来很简单，也没有什么特别的，是吗? 接下来我们要对比几种不同的记忆机制，发现他们的不同用途。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 摘要对话记忆 ConversationSummaryMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.conversation.memory import ConversationSummaryMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_memory = ConversationSummaryMemory(llm=OpenAI(openai_api_key=OPENAI_API_TOKEN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    memory=summary_memory\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi there! I am Sam\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\""
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "The human introduces themselves to the AI. The AI responds with a greeting and asks what brings the human to them.\n",
      "Human: How are you today?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" I'm doing great, thank you for asking! What brings you to me today?\""
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"How are you today?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "The human introduces themselves to the AI and asks how they are doing. The AI responds positively and asks what brings the human to them.\n",
      "Human: Can you help me with some customer support?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Absolutely! I'm here to help. What kind of customer support do you need?\""
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Can you help me with some customer support?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "区别在于 `Current conversation:`的内容。 运行后可以发现，AI在总结对话：“The human introduces themselves to the AI as Sam and the AI introduces itself. ”\n",
    "\n",
    "而\"what brings Sam here today?\"和 \"How are you today?\"这样的句子都被省略了。\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们打印看看 `conversation.memory.buffer` 里存了一些什么内容。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The human introduces themselves to the AI and asks how they are doing. The AI responds positively and asks what brings the human to them. The human then asks for customer support and the AI responds positively, offering to help and asking what kind of support is needed.\n"
     ]
    }
   ],
   "source": [
    "print(conversation.memory.buffer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "很有意思对吧，AI自动对对话进行了摘要。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 对话缓冲窗口记忆 ConversationBufferWindowMemory"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "还是像之前一样，导入和实例化 `ConversationBufferWindowMemory`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.conversation.memory import ConversationBufferWindowMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "window_memory = ConversationBufferWindowMemory(llm=OpenAI(openai_api_key=OPENAI_API_TOKEN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    memory=window_memory\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation.predict(input=\"I am looking for some customer support\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: I am looking for some customer support\n",
      "AI:  Sure thing! I can help you with that. What kind of customer support do you need?\n",
      "Human: My TV is not working.\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" I'm sorry to hear that. Can you tell me more about the issue you're having with your TV?\""
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"My TV is not working.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以看到 `Current conversation:` 一直在把 AI 和 Human 的对话保存起来。但是如果我们改变交互的次数，会发生什么？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We set a low k=2, to only keep the last 2 interactions in memory\n",
    "window_memory = ConversationBufferWindowMemory(k=2,llm=OpenAI(openai_api_key=OPENAI_API_TOKEN))\n",
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    memory=window_memory\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi there! I am Sam\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\""
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: I am looking for some customer support\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Sure thing! I can help you with that. What kind of customer support do you need?'"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"I am looking for some customer support\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there! I am Sam\n",
      "AI:  Hi Sam! My name is AI. It's nice to meet you. What brings you here today?\n",
      "Human: I am looking for some customer support\n",
      "AI:  Sure thing! I can help you with that. What kind of customer support do you need?\n",
      "Human: When I turn it on it makes some weird sounds and then goes black\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hmm, that doesn't sound good. What kind of device are you having trouble with?\""
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"When I turn it on it makes some weird sounds and then goes black\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: I am looking for some customer support\n",
      "AI:  Sure thing! I can help you with that. What kind of customer support do you need?\n",
      "Human: When I turn it on it makes some weird sounds and then goes black\n",
      "AI:  Hmm, that doesn't sound good. What kind of device are you having trouble with?\n",
      "Human: My TV is not working.\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" I see. That's definitely not ideal. Can you tell me a bit more about the issue? What kind of weird sounds does it make when you turn it on?\""
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"My TV is not working.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以注意到 Current conversation 现在只保留了最后2次的交互对话。而其他的对话都省略了。当我们设置 K = 2 的时候，限制了记忆存储的对话次数。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 对话摘要缓冲记忆 ConversationSummaryBufferMemory"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这个看起来是把最近的交互存储起来，但是做了摘要。我们来看看究竟是如何？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.conversation.memory import ConversationSummaryBufferMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model_name='text-davinci-003', \n",
    "             temperature=0, \n",
    "             max_tokens = 512,\n",
    "             openai_api_key=OPENAI_API_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# max_token_limit=40 - token limits needs transformers installed\n",
    "memory = ConversationSummaryBufferMemory(llm=llm,max_token_limit=40) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm, \n",
    "    memory=memory, \n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们还是用之前的对话内容，来测试。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_summary.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_summary.predict(input=\"I need help with my broken TV\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_summary.predict(input=\"It makes weird sounds when i turn it on and then goes black\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_summary.predict(input=\"It seems to be Hardware\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后打印一下，最后的对话记忆有什么内容？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(conversation_with_summary.memory.moving_summary_buffer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "\n",
    "The human, Sam, introduces themselves to the AI. The AI responds and asks what brings Sam to them today, and Sam responds that they need help with their broken TV. The AI expresses sympathy and asks what the problem with the TV is, and Sam responds that it makes weird sounds when they turn it on and then goes black. The AI responds with \"Hmm, that doesn't sound good. Do you know if it's a hardware or software issue?\"\n",
    "\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 知识图谱记忆 ConversationKGMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.conversation.memory import ConversationKGMemory\n",
    "from langchain.prompts.prompt import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model_name='text-davinci-003', \n",
    "             temperature=0, \n",
    "             max_tokens = 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. \n",
    "If the AI does not know the answer to a question, it truthfully says it does not know. The AI ONLY uses information contained in the \"Relevant Information\" section and does not hallucinate.\n",
    "\n",
    "Relevant Information:\n",
    "\n",
    "{history}\n",
    "\n",
    "Conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"history\", \"input\"], template=template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_kg = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    prompt=prompt,\n",
    "    memory=ConversationKGMemory(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_kg.predict(input=\"Hi there! I am Sam\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_kg.predict(input=\"My TV is broken and I need some customer assistance\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_kg.predict(input=\"It makes weird sounds when i turn it on and then goes black\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_with_kg.predict(input=\"Yes it is and it is still under warranty. my warranty number is A512423\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NetworkX Entity graph\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# nx.draw(conversation_with_kg.memory.kg, with_labels=True)\n",
    "# plt.show()\n",
    "\n",
    "print(conversation_with_kg.memory.kg)\n",
    "print(conversation_with_kg.memory.kg.get_triples())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以看到知识图谱的记忆内容是：\n",
    "\n",
    "```\n",
    "<langchain.graphs.networkx_graph.NetworkxEntityGraph object at 0x7f2bc65fe5e0>\n",
    "[('Sam', 'Human', 'is'), ('TV', 'true', 'is broken'), ('TV', 'weird sounds', 'makes'), ('TV', 'black', 'goes'), ('TV', 'yes', 'is under warranty'), ('device', 'A512423', 'is under warranty')]\n",
    "\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 实体记忆 ConversationEntityMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import OpenAI, ConversationChain\n",
    "from langchain.chains.conversation.memory import ConversationEntityMemory\n",
    "from langchain.chains.conversation.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE\n",
    "from pydantic import BaseModel\n",
    "from typing import List, Dict, Any"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们打印看看 Langchain 的 实体记忆 prompt 提示模板写了些什么？\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'You are an assistant to a human, powered by a large language model trained by OpenAI.\\n\\nYou are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\\n\\nYou are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\\n\\nOverall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\\n\\nContext:\\n{entities}\\n\\nCurrent conversation:\\n{history}\\nLast line:\\nHuman: {input}\\nYou:'"
      ]
     },
     "execution_count": 123,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "####   The propmpt\n",
    "ENTITY_MEMORY_CONVERSATION_TEMPLATE.template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model_name='text-davinci-003', \n",
    "             temperature=0, \n",
    "             max_tokens = 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True,\n",
    "    prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,\n",
    "    memory=ConversationEntityMemory(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation.predict(input=\"Hi I am Sam. My TV is broken but it is under warranty.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation.predict(input=\"How can I get it fixed. The warranty number is A512453\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation.memory.entity_cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "pprint(conversation.memory.store)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "{'A512453': \"A512453 is the warranty number for Sam's TV.\",\n",
    " 'Dave': 'Dave is a repair person.',\n",
    " 'Sam': 'Sam owns a TV that is currently broken and under warranty.',\n",
    " 'TV': 'TV is under warranty.'}\n",
    " ```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
